package com.yys.apm.service;

import com.yys.apm.boot.AgentBootLoader;
import com.yys.apm.boot.ApmAgentManager;
import com.yys.apm.config.ApmConfig;
import com.yys.apm.boot.spi.CollectSPI;
import com.yys.apm.boot.spi.ServiceSPI;
import com.yys.apm.collect.ApmCollect;
import com.yys.apm.logger.ILogger;
import com.yys.apm.logger.LoggerManager;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.LoaderClassPath;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 采集服务
 *
 * @author Yingsheng Ye
 * @date 2021-03-19 18:17
 */
@ServiceSPI
public class CollectService implements ApmService {

    private static final ILogger logger = LoggerManager.getLogger(CollectService.class);

    private static final String SPLIT_DOT = ",";

    private Properties properties;

    private Instrumentation instrumentation;

    private Map<ClassLoader, ClassPool> classPoolMap = new ConcurrentHashMap<>();

    private Map<Class, ApmCollect> collects = Collections.emptyMap();

    @Override
    public void prepare() {
        ApmAgentManager apmAgentManager = ApmAgentManager.getInstance();
        properties = apmAgentManager.properties;
        instrumentation = apmAgentManager.instrumentation;
        collects = loadServices();
    }

    @Override
    public void startup() {
        instrumentation.addTransformer(new DefineApmClassFileTransformer());
    }

    @Override
    public void onComplete() {

    }

    @Override
    public void shutdown() {

    }

    private Map<Class, ApmCollect> loadServices() {
        Map<Class, ApmCollect> services = new LinkedHashMap<>(16);
        for (ApmCollect service : ServiceLoader.load(ApmCollect.class)) {
            Class<? extends ApmCollect> serviceClass = service.getClass();
            boolean isCollect = serviceClass.isAnnotationPresent(CollectSPI.class);
            if (isCollect && !services.containsKey(serviceClass)) {
                services.put(serviceClass, service);
            }
        }
        return services;
    }

    class DefineApmClassFileTransformer implements ClassFileTransformer {
        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classFileBuffer) {
            if (className == null || loader == null || className.contains("$$")) {
                return new byte[0];
            }
            className = className.replaceAll("/", ".");
            String scanBasePackage = properties.getProperty(ApmConfig.SCAN_BASE_PACKAGE);
            if (null != scanBasePackage && !"".equals(scanBasePackage)) {
                String[] scanBasePackages = scanBasePackage.split(SPLIT_DOT);
                boolean flag = false;
                for (String pkg : scanBasePackages) {
                    if (className.trim().startsWith(pkg.trim())) {
                        flag = true;
                        break;
                    }
                }
                if (!flag) {
                    return new byte[0];
                }
            }
            if (!classPoolMap.containsKey(loader)) {
                ClassPool classPool = new ClassPool();
                classPool.insertClassPath(new LoaderClassPath(loader));
                classPoolMap.put(loader, classPool);
            }
            ClassPool cp = classPoolMap.get(loader);
            try {
                CtClass ctClass = cp.get(className);
                for (ApmCollect c : collects.values()) {
                    AgentBootLoader agentLoader = new AgentBootLoader(loader, className, classFileBuffer, ctClass);
                    if (c.isApmTarget(agentLoader)) {
                        return c.transform(agentLoader);
                    }
                }
            } catch (Exception e) {
                logger.error(e,"%s first-apm agent fail", className);
            }
            return new byte[0];
        }
    }

}
